from random import random
from elementtree.ElementTree import parse
from readGear import *
from math import ceil
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
import psyco

psyco.full()

def roll(num):
   result = 0
   for i in range(num):
      die = ceil(random() * 6)
      if die == 6 and result >= 6:
         result += 1
      elif die > result:
         result = die
   if result == 1:
      result = -10
   return int(result)

def directSim(atkDice, defDice, atkMod, defMod):
   atkRoll = roll(atkDice) + atkMod
   if atkRoll < 0:
      atkRoll = 0
   defRoll = roll(defDice) + defMod
   if defRoll < 0:
      defRoll = 0
   mos = atkRoll - defRoll
   return mos

def indirectSim(atkDice, defDice, atkMod, defMod, ae, sensors, thrown):
   atkRoll = roll(atkDice)
   mos = atkRoll + atkMod
   mos -= 4
   if mos > 0:
      hit = 1
   else:
      scatter = randint(1, 6) - mos + sensors
      if thrown:
         scatter = scatter / 2
      if ae >= scatter:
         hit = 2
      else:
         hit = 0
   if hit:
      defRoll = roll(defDice) + defMod
      if defRoll < 0:
         defRoll = 0
      mos = atkRoll - defRoll 
   else:
      mos = -1
   return mos, hit

def damage(dm, mos, armor):
   dam = dm * mos / armor
   if dam < 0:
      dam = 0
   if dam >= 3:
      dam = 10
   return dam

def fire(attacker, defender, distance, weapon, indirect=0):
   weapon = attacker.weapons[weapon]
   atkMod = attacker.getPenelty() + 1
   defDice = defender.pilot['def']
   if weapon.ai:
      if not defender.perks.has_key('Infantry'):
         defDice += 1
   defMod = defender.getPenelty()
   if defender.ca['stunned']:
      if defender.getDodge() > 0:
         defMod -= 1
         
   dm = weapon.dm
   rof = weapon.rof
   spray = 0
   burst = 0
   ae = 0
   hit = 0
   armor = defender.armor
   if defender.perks.has_key('reinforced armor front'):
      armor += defender.perks['reinforced armor front']

   if weapon.melee:
      atkDice = attacker.pilot['def']
      for i in defender.weapons:
         if defender.weapons[i].melee:
            if defender.weapons[i].acc > 0:
               defMod += 1
               break
      if attacker.ca['speedband'] > defender.ca['speedband']:
         atkMod += 1
      if attacker.size > defender.size: # Size bonuses.
         atkMod += 1
      elif defender.size > attacker.size:
         defMod += 1
      if weapon.code == 'VK':
         if attacker.weapons.has_key('VB'):
            atkDice += 1
      mos = directSim(atkDice, defDice, atkMod, defMod)
   else: # Ranged combat!
      atkDice = attacker.pilot['atk']
      if attacker.perks.has_key('defective fire control'):
         atkDice -= attacker.perk['defective fire control']
      defMod += defender.getDodge()
      rangeband = weapon.getRangeband(distance)
      if weapon.notes.has_key('mb'):
         mb = (rangeband - 2) * weapon.notes['mb']
         if mb < 0:
            mb = 0
         dm -= mb
      atkMod += attacker.getAttackMod() + weapon.getRangeMod(rangeband)

      if indirect or weapon.thrown:
         if weapon.ae:
            if weapon.ai:
               ae += weapon.ae
            else:
               # This decides if the attack will use AE. If it can stun and the stun will lower
               # def, it will always try to stun. It also decides if to use RoF for spray(AE) or burst(DM)
               # and in what amounts. It'll try and hit the dm=armor breakpoint if it can, and if not,
               # it'll just go for spray.
               if defender.getDodge() > 0 and defender.ca['stunned'] == 0:
                  if (dm - weapon.ae) * 2 > armor:
                     dm -= weapon.ae
                     ae += weapon.ae
                  elif (dm - weapon.ae + rof) * 2 > armor:
                     spray = (((dm - weapon.ae + rof) * 2) - armor) / 2
                     burst = weapon.rof - spray
                     rof = 0
                     dm -= weapon.ae
                     ae += weapon.ae
               else:
                  if dm - weapon.ae > armor:
                     ae += weapon.ae
                     dm -= weapon.ae
                  elif dm + rof - weapon.ae > armor:
                     spray = dm - weaponae + rof - armor
                     burst = weapon.rof - spray
                     rof = 0
                     ae += weapon.ae
                     dm -= weapon.ae
         if rof:
            if dm >= armor:
               spray = rof
               rof = 0
            elif dm + rof >= armor:
               spray = dm + rof - armor
               burst = weapon.rof - spray
               rof = 0
            else:
               spray = rof
               rof = 0
         if weapon.thrown:
            mos = indirectSim(atkDice, defDice, atkMod, defMod, ae + spray, attacker.sensors, 1)
         else:
            defMod -= 1
            mos = indirectSim(atkDice, defDice, atkMod, defMod, ae + spray, attacker.sensors, 0)

         if weapon.ae:
            if mos[1]:
               hit = 1            
         mos = mos[0]     
      else:
         if weapon.ae:
            if defender.getDodge() > 0 and defender.ca['stunned'] == 0:
               ae = weapon.ae
               dm -= weapon.ae
         dm += rof
         rof = 0
         mos = directSim(atkDice, defDice, atkMod, defMod)
         
   if weapon.notes.has_key('ap'):
      if mos >= 0:
         mos += 2

   if ae:
      if mos > 0 or hit:
         if dm * 2 > armor:
            defender.ca['stunned'] = 1

   dam = damage(mos, dm, armor)

   if weapon.notes.has_key('h'):
      die = ceil(random() * 6)
      if defender.perks.has_key('vuln to haywire'):
         if die >= 3:
            dam += 1
         else:
            defender.ca['stunned'] = 1
      else:
         if die == 6:
            dam += 1
         elif die >= 3:
            defender.ca['stunned'] = 1

   return dam

def testWeapons(attacker, defender, distance, weapons):
   total = 0
   for i in xrange(1, 1001):
      turns = 0
      defender.ca['damage'] = 0
      defender.ca['stunned'] = 0
      while not defender.isDead():
         turns += 1
         total += 1
         if turns >= 20:
            break
         for weapon in weapons:
            if weapon.find(' IF') > -1:
               defender.ca['damage'] += fire(attacker, defender, distance, weapon.split(' IF')[0], indirect=1)
            else:
               defender.ca['damage'] += fire(attacker, defender, distance, weapon)
   return total / 1000.0

def plotGear(attacker, atkPilot, defender, defPilot, options=0, swaps=0):
   option = ''
   if options:
      for i in options:
         option += i + '/'
   option = option[:-1]
   fig = Figure()
   canvas = FigureCanvas(fig)
   ax = fig.add_subplot(111)
   ax.set_title('%s %s(%s/%s) - %s(%s/%s)' % (option, attacker, atkPilot['atk'], atkPilot['def'], defender, defPilot['atk'], defPilot['def']))
   ax.set_ylabel('Shots to kill')
   ax.set_xlabel('Distance')
   ax.set_yticks(range(0, 21, 2))
   ax.set_xticks(range(0, 25, 2))
   ax.grid (True)
   
   attacker = Gear(attacker, atkPilot)
   if options:
      for option in options:
         attacker.addWeapon(option)
   if swaps:
      for swap in swaps:
         attacker.removeWeapon(swap)
   defender = Gear(defender, defPilot)
   if defender.movement.has_key('Ground'):
      defender.ca['movementMode'] = 'Ground'

   # Tests each weapon the attacker has at all of its rangebands.
   results = {}
   for weapon in attacker.weapons:
      results[weapon] = []
      w = attacker.weapons[weapon]
      for distance in w.rng:
         score = testWeapons(attacker, defender, distance, [weapon])
         if w.indirect:
            indirectScore = testWeapons(attacker, defender, distance, [weapon])
            if indirectScore > score:
               score = indirectScore
         results[weapon].append(score)

   if attacker.actions > 1:
      # Multiaction units are still done bluntly but it's less important as they kill things faster.
      # This sorts the unit's weapons and test the top (number of action) ones at all ranges 1-24.
      # If a unit's best weapon if an IF weapon, it will use it directly instead.
      lowestResults = []
      for distance in xrange(1, 25):
         bestWeapons = {}
         for weapon in results:
            for rangeband in attacker.weapons[weapon].rng:
               if distance <= rangeband:
                  bestWeapons[weapon] = results[weapon][attacker.weapons[weapon].getRangeband(rangeband)-1]
         items = bestWeapons.items()
         backitems = [[v[1], v[0]] for v in items]
         backitems.sort()
         sortedWeapons = [backitems[i][1] for i in range(0, len(backitems))]
         lowestResults.append(testWeapons(attacker, defender, distance, sortedWeapons[:attacker.actions]))
   else:
      # This converts the breakpoint scores of results into a full list of the best scores from ranges 1-24
      lowestResults = range(20, 44)
      for distance in xrange(1, 25):
         for weapon in results:
            for rangeband in attacker.weapons[weapon].rng:
               if distance <= rangeband:
                  if results[weapon][attacker.weapons[weapon].getRangeband(rangeband)-1] < lowestResults[distance-1]:
                     lowestResults[distance-1] = results[weapon][attacker.weapons[weapon].getRangeband(rangeband)-1]
                     
   ax.plot(xrange(1, 25), lowestResults)
   ax.axis([1, 24, 0, 20])
   canvas.print_figure('%s %s(%s%s) - %s(%s%s).png' % (option, attacker.name, atkPilot['atk'], atkPilot['def'], defender.name, defPilot['atk'], defPilot['def']))
   return results

atkGears = ['Bear', 'Mauler Bear', 'Black Cat', 'Bobcat', 'Cheetah Paratrooper',
           'Cheetah', 'Ferret', 'Ferret MKII', 'Grizzly', 'Hunter Commando',
           'Hunter Paratrooper', 'Hunter', 'Jaguar', 'Kodiak', 'Mad Dog R',
           'Panther', 'Rabid Grizzly', 'Razorback', 'Peacemaker Razorback',
           'Strike Cheetah', 'Stripped-Down Hunter', 'Tiger', 'Aller',
           'Bandit Hunter Klemm', 'Hardy Aller', 'Jaxon Support Tank', 'Klemm',
           'Naval Support Aller', 'Mammoth', 'Thunderhammer', 'Anolis Reit',
           'Black Adder', 'Black Mamba MP', 'Cobra MP', 'Fire Dragon',
           'Iguana Commando', 'Iguana MP', 'Iguana', 'King Cobra', 'Naga',
           'Python', 'Rattlesnake', 'Sidewinder', 'Snakeeye Black Mamba',
           'Spitting Cobra', 'Hetairoi', 'Hun', 'Visigoth']

defGears = ('Cheetah', 'Anolis Refit', 'Dartjager', 'Ferret', 'Strike Cheetah',
           'Rattlesnake', 'Bobcat', 'Silverscale', 'Iguana', 'Hunter', 'Stripped-Down Hunter',
           'Hunter Commando', 'Mad Dog R', 'Black Adder', 'Sidewinder', 'Jaguar',
           'Tiger', 'Black Mamba', 'Bear', 'Grizzly', 'Python', 'Spitting Cobra', 'King Cobra',
           'Kodiak', 'Naga', 'Mammoth', 'Hun', 'Aller')

defaultPilot = {'atk':2, 'def':2, 'ew':1, 'ld':1}
vetPilot = {'atk':3, 'def':3, 'ew':1, 'ld':1}
acPilot = {'atk':4, 'def':4, 'ew':1, 'ld':1}

pilots = defaultPilot, vetPilot, acPilot

##for defender in defGears:
##   plotGear('Jaguar', pilot, defender, pilot)
##   print defender
